use std::sync::Arc;

use async_trait::async_trait;
use datafusion::dataframe::DataFrame;
use datafusion::datasource::TableProvider;
use datafusion::prelude::{col, Expr};
use log::{debug, warn};
use parking_lot::RwLock;

use common::{Entry, Error, LResult, SchemaMut};
use engine::{Engine, EngineFactory, EngineFactoryRef, EngineRef};
use function::call_symbol_to_string;

use crate::mem_table::MemTable;
use crate::table_provider::MemTableProvider;

pub struct MemEngine {
    measurement: String,
    mem_table: Arc<RwLock<MemTable>>,
}

pub struct MemEngineFactory;

impl MemEngine {
    pub fn create<S: Into<String>>(name: S) -> Self {
        let measurement = name.into();
        let inner = MemTable::create(measurement.clone());
        let mem_table = Arc::new(RwLock::new(inner));
        Self {
            measurement,
            mem_table,
        }
    }
}

#[async_trait]
impl Engine for MemEngine {
    #[inline]
    fn name(&self) -> &String {
        &self.measurement
    }
    async fn write(&mut self, entry: Entry) -> LResult<()> {
        self.mem_table.write().push(entry)
    }

    fn table_provider(&self) -> LResult<Arc<dyn TableProvider>> {
        Ok(Arc::new(MemTableProvider::create(self.mem_table.clone())?))
    }

    async fn delete(&mut self) -> LResult<()> {
        warn!("Table {} deleted", self.name());
        Ok(())
    }
}

impl MemEngineFactory {
    #[inline]
    pub fn create() -> EngineFactoryRef {
        Box::new(Self {})
    }
}

impl EngineFactory for MemEngineFactory {
    fn name(&self) -> &'static str {
        static NAME: &'static str = "mem_table";
        &NAME
    }

    fn create(&self, measurement: String) -> EngineRef {
        Box::new(MemEngine::create(measurement))
    }

    fn post_transform_query(&self, df: DataFrame) -> LResult<DataFrame> {
        let mut expr_list: Vec<Expr> = vec![];
        // transform all tag(u64) to tag(Utf8)
        for f in df.schema().fields() {
            if SchemaMut::is_real_tag(f.name()) {
                // project column with symbol_to_string() function
                expr_list.push(call_symbol_to_string(f.name()));
            } else {
                // direct project this column
                expr_list.push(col(f.name()))
            }
        }
        #[cfg(debug_assertions)]
        {
            debug!(target: "engine:mem", "transform_query select {:?}", expr_list);
        }
        df.select(expr_list).map_err(|e| {
            Error::ExecutionErr(format!(
                "transform_query err: {:?}", e
            ))
        })
    }
}